package com.zyb.lmbackend.util;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * 支持解析 LLM 输出为 单个 JSON 数组 或 JSON Lines（每行一个 JSON 对象）。
 * - 去除 UTF-8 BOM、Markdown 代码块围栏、首尾空白
 * - JSON 数组：整体为一个数组
 * - JSON Lines：按行解析为多个对象
 */
/**
 * 将 LLM 文本输出解析为 JSON 节点序列。
 *
 * 支持两种格式：
 * - 单个 JSON 数组（整体为 [ ... ]）
 * - JSON Lines（按行一个 JSON 对象）
 * 并处理 UTF-8 BOM 与 Markdown 代码块围栏。
 */
public class JsonLinesOrArrayParser {

    private final ObjectMapper mapper;

    public JsonLinesOrArrayParser(ObjectMapper mapper) {
        this.mapper = mapper;
    }

    /**
     * 解析原始字符串为 JSON 节点列表。
     * @throws InvalidLlmOutputException 当为空或解析失败
     */
    public List<JsonNode> parseToNodes(String raw) {
        if (raw == null) throw new InvalidLlmOutputException("LLM 输出为空");
        String s = sanitize(raw);
        if (s.isBlank()) throw new InvalidLlmOutputException("LLM 输出为空白");

        // 判断是否为数组
        if (s.startsWith("[") && s.endsWith("]")) {
            try {
                JsonNode node = mapper.readTree(s);
                if (!node.isArray()) {
                    throw new InvalidLlmOutputException("期望为 JSON 数组，但解析到的根节点不是数组");
                }
                List<JsonNode> list = new ArrayList<>();
                node.forEach(list::add);
                return list;
            } catch (JsonProcessingException e) {
                throw new InvalidLlmOutputException("JSON 数组解析失败", e);
            }
        }

        // 逐行 JSON 对象
        String[] lines = s.split("\r?\n");
        List<JsonNode> list = new ArrayList<>();
        int lineNo = 0;
        for (String line : lines) {
            lineNo++;
            String t = line.trim();
            if (t.isEmpty()) continue;
            try {
                JsonNode node = mapper.readTree(t);
                if (!node.isObject()) {
                    throw new InvalidLlmOutputException("第 " + lineNo + " 行不是 JSON 对象");
                }
                list.add(node);
            } catch (JsonProcessingException e) {
                throw new InvalidLlmOutputException("第 " + lineNo + " 行 JSON 解析失败: " + t, e);
            }
        }
        if (list.isEmpty()) {
            throw new InvalidLlmOutputException("未解析到任何 JSON 对象");
        }
        return list;
    }

    /** 对输入文本做清洗：去 BOM、去围栏、去首尾空白。 */
    private String sanitize(String s) {
        // 去除 BOM
        byte[] b = s.getBytes(StandardCharsets.UTF_8);
        if (b.length >= 3 && b[0] == (byte) 0xEF && b[1] == (byte) 0xBB && b[2] == (byte) 0xBF) {
            s = new String(b, 3, b.length - 3, StandardCharsets.UTF_8);
        }
        s = s.trim();
        // 去除 Markdown 围栏
        if (s.startsWith("```")) {
            s = s.replaceAll("(?s)^```[a-zA-Z0-9]*\\s*", "");
            s = s.replaceAll("(?s)```\\s*$", "");
            s = s.trim();
        }
        return s;
    }
}
